import rsa
import hashlib
import base64
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad

class RsaEncryption(object):
    def __init__(self) -> None:
        super().__init__()
        self._cipher = None
        self._pubkey = None
        self._privkey = None

    def aes_set_key(self, key:bytes):
        self._cipher = AES.new(key, AES.MODE_ECB)

    def aes_encrypt(self, data:bytes)->bytes:
        try:
            return self._cipher.encrypt(pad(data, 16)) 
        except Exception as e:
            print(f'aes encrypt failed {e}')
            return None

    def aes_decryypt(self, data:bytes)->bytes:
        try:
            return unpad(self._cipher.decrypt(data), 16)
        except Exception as e:
            print(f'aes decryypt failed {e}')
            return None

    def rsa_set_pem_file(self, pub_pem:str, pri_pem:str):
        try:
            with open(pub_pem, 'rb') as f:
                self._pubkey = rsa.PublicKey.load_pkcs1(f.read())
            with open(pri_pem, 'rb') as f:
                self._privkey = rsa.PrivateKey.load_pkcs1(f.read())
        except Exception as e:
            print(f'rsa_set_pem_file failed {e}')
            return None
    
    def rsa_encrypt(self, data:bytes)->bytes:
        try:
            return rsa.encrypt(data, self._pubkey)
        except Exception as e:
            print(f'rsa_encrypt failed {e}')
            return None

    def rsa_decrypt(self, data:bytes)->bytes:
        try:
            return rsa.decrypt(data, self._privkey)
        except Exception as e:
            print(f'rsa_decrypt failed {e}')
            return None


# if __name__ == '__main__':
    
#     from uuid import uuid4
#     key = uuid4().hex
#     data = 'liubenlong'.encode()
#     enc = RsaEncryption()
#     dec = RsaEncryption()
#     if 0:
#         iv = key[0:16].encode()
#         enc.aes_set_key(key.encode(), iv)
#         dec.aes_set_key(key.encode(), iv)
#         enc_data = enc.aes_encrypt(data)
#         print(f'data is {data} and aes enc data is {enc_data} len {len(enc_data)}')
#         dec_data = dec.aes_decryypt(enc_data)
#         print(f'dec data is {dec_data}')
#     else:
#         enc.rsa_set_pem_file('public.pem','private.pem')

#         enc_data = enc.rsa_encrypt(data)
#         print(f'rsa enc_data {enc_data} data {data}')
#         dec_data = enc.rsa_decrypt(enc_data)
#         print(f'rsa dec data {dec_data}')

    

